package com.blyat.xsoft.kernel.component.mq.rocketmq.consumer.pull;

import com.alibaba.fastjson.JSON;
import com.blyat.xsoft.kernel.component.mq.MQConsumer;
import com.blyat.xsoft.kernel.component.mq.rocketmq.Carrier;
import com.blyat.xsoft.kernel.component.mq.rocketmq.RocketMQContextLoader;
import com.google.common.collect.Lists;
import org.apache.rocketmq.client.consumer.DefaultLitePullConsumer;
import org.apache.rocketmq.common.consumer.ConsumeFromWhere;
import org.apache.rocketmq.common.message.MessageExt;
import org.apache.rocketmq.common.protocol.heartbeat.MessageModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.util.StopWatch;

import java.util.List;
import java.util.Objects;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @Auther: syh
 * @Date: 2020/8/19
 * @Description:
 */
public abstract class PullConsumer<T> extends RocketMQContextLoader
        implements MQConsumer<T>, DisposableBean {
    private final static Logger log = LoggerFactory.getLogger(PullConsumer.class);

    private final static String DEFAULT_TAGS = "*";
    private final static int DEFAULT_BATCH_SIZE = 32;
    private final static int DEFAULT_PULL_THREAD_NUMS = 20;
    private final static long DEFAULT_POLL_TIMEOUT_MILLIS = 3 * 1000l;
    private final int DEFAULT_IO_THREADS = Math.min(Runtime.getRuntime().availableProcessors() + 1, 32);
    private final int DEFAULT_KEEP_ALIVE_TIME = 60 * 1000;
    private final int DEFAULT_BLOCKING_QUEUE_SIZE = 2000;

    private String groupName;
    private String topic;
    private String tags;
    private int batchSize;
    private int pullThreadNums;
    private long pollTimeoutMillis;

    private ExecutorService executorService;
    private DefaultLitePullConsumer consumer;
    private RocketMqConsumerRunner rocketMqConsumerRunner;

    public PullConsumer(String topic) {
        this(topic, topic, DEFAULT_TAGS);
    }

    public PullConsumer(String groupName, String topic) {
        this(groupName, topic, DEFAULT_TAGS);
    }

    public PullConsumer(String groupName, String topic, String tags) {
        this(topic, tags, DEFAULT_BATCH_SIZE);
    }

    public PullConsumer(String topic, int batchSize) {
        this(topic, DEFAULT_TAGS, batchSize);
    }

    public PullConsumer(String topic, String tags, int batchSize) {
        this(topic, topic, tags, batchSize);
    }

    public PullConsumer(String groupName, String topic, String tags, int batchSize) {
        this(groupName, topic, tags, batchSize, DEFAULT_PULL_THREAD_NUMS, DEFAULT_POLL_TIMEOUT_MILLIS);
    }

    public PullConsumer(String groupName, String topic, String tags, int batchSize, int pullThreadNums, long pollTimeoutMillis) {
        this.groupName = groupName;
        this.topic = topic;
        this.tags = tags;
        this.batchSize = batchSize;
        this.rocketMqConsumerRunner = new RocketMqConsumerRunner();
        this.pullThreadNums = pullThreadNums;
        this.pollTimeoutMillis = pollTimeoutMillis;
        this.executorService = new ThreadPoolExecutor(1, DEFAULT_IO_THREADS,
                DEFAULT_KEEP_ALIVE_TIME, TimeUnit.MILLISECONDS, new LinkedBlockingDeque<>(DEFAULT_BLOCKING_QUEUE_SIZE));
    }

    @Override
    protected void doAfterSingletonsInstantiated() {
        createRocketMQConsumer();
    }

    private void createRocketMQConsumer() {
        try {
            consumer = new DefaultLitePullConsumer(groupName);
            consumer.setNamesrvAddr(getProperties().getNameServer());

            consumer.setPullBatchSize(batchSize);
            consumer.setPullThreadNums(pullThreadNums);
            consumer.setPollTimeoutMillis(pollTimeoutMillis);
            /*
             * CONSUME_FROM_LAST_OFFSET 默认策略，从该队列最尾开始消费，跳过历史消息
             * CONSUME_FROM_FIRST_OFFSET 从队列最开始开始消费，即历史消息（还储存在broker的）全部消费一遍
             */
            consumer.setConsumeFromWhere(ConsumeFromWhere.CONSUME_FROM_FIRST_OFFSET);
            // 默认是集群方式，可以更改为广播，但是广播方式不支持重试
            consumer.setMessageModel(MessageModel.CLUSTERING);

            consumer.subscribe(topic, tags);
            consumer.start();
            executorService.execute(rocketMqConsumerRunner);
        } catch (Exception e) {
            //
        }
    }

    @Override
    public void destroy() throws Exception {
        if (this.consumer != null) {
            this.consumer.shutdown();
        }
    }

    private class RocketMqConsumerRunner implements Runnable {
        private List<T> data;
        private boolean runnable;

        public RocketMqConsumerRunner() {
            this.data = Lists.newArrayListWithExpectedSize(PullConsumer.this.batchSize);
            this.runnable = true;
        }

        @Override
        public void run() {
            while (runnable) {
                try {
                    List<MessageExt> messageList = PullConsumer.this.consumer.poll();
                    if (Objects.isNull(messageList) || messageList.isEmpty()) {
                        // flush the rest data.
                        flush();
                        continue;
                    }

                    List<T> list = messageList.parallelStream()
                            .map(message -> {
                                Carrier<T> carrier = JSON.parseObject(new String(message.getBody()), Carrier.class);
                                return carrier.getContent();
                            }).collect(Collectors.toList());

                    if (data.size() >= PullConsumer.this.batchSize) {
                        flush();
                    }
                    data.addAll(list);
                    // 不执行该方法，应用重启会存在重复消费。
                    PullConsumer.this.consumer.commitSync();
                } catch (Exception e) {
                    log.error("rocketmq pull consume failed.", e);
                }
            }
        }

        private void flush() throws Exception {
            if (Objects.isNull(data) || data.isEmpty()) return;
            StopWatch watch = new StopWatch();
            watch.start();
            PullConsumer.this.consume(data);
            watch.stop();
            log.info("flush data size: {} used times:{}", data.size(), watch.getLastTaskTimeMillis());
            data.forEach(d -> d = null);
            data.clear();
        }

        public void setRunnable(boolean runnable) {
            this.runnable = runnable;
        }
    }
}
